import subprocess
import sys
import os
import signal
import time
from datetime import datetime

# 颜色代码
class Colors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

def log(message: str, level: str = "INFO"):
    """打印带时间戳和颜色的日志"""
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

    # 根据日志级别选择颜色
    color_map = {
        "INFO": Colors.OKBLUE,
        "SUCCESS": Colors.OKGREEN,
        "WARNING": Colors.WARNING,
        "ERROR": Colors.FAIL,
        "HEADER": Colors.HEADER
    }

    color = color_map.get(level, Colors.OKBLUE)
    print(f'{color}[{timestamp}] [{level}] {message}{Colors.ENDC}')

def print_banner():
    """打印启动横幅"""
    banner = f"""
{Colors.HEADER}{'='*60}
    煤炉监控系统 - 统一启动脚本
    FastAPI + Celery + Redis + MySQL
    Version: 1.0.0
{'='*60}{Colors.ENDC}
"""
    print(banner)

def check_dependencies():
    """检查系统依赖"""
    log("检查系统依赖...", "INFO")

    # 检查Python版本
    python_version = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
    log(f"Python版本: {python_version}", "INFO")

    # 检查必要的模块
    required_modules = ['uvicorn', 'celery', 'fastapi', 'sqlalchemy', 'redis']
    missing_modules = []

    for module in required_modules:
        try:
            __import__(module)
            log(f"[OK] {module} 模块已安装", "SUCCESS")
        except ImportError:
            missing_modules.append(module)
            log(f"[FAIL] {module} 模块未安装", "ERROR")

    if missing_modules:
        log(f"缺少依赖模块: {', '.join(missing_modules)}", "ERROR")
        log("请运行: pip install -r requirements.txt", "WARNING")
        return False

    log("所有依赖检查通过", "SUCCESS")
    return True

def check_ports():
    """检查端口占用情况"""
    log("检查端口占用情况...", "INFO")

    import socket
    ports_to_check = [8000, 6379, 5672]  # FastAPI, Redis, RabbitMQ

    for port in ports_to_check:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            result = sock.connect_ex(('localhost', port))
            if result == 0:
                log(f"端口 {port} 已被占用", "WARNING")
            else:
                log(f"端口 {port} 可用", "SUCCESS")
        except Exception as e:
            log(f"检查端口 {port} 时出错: {e}", "ERROR")
        finally:
            sock.close()

def start_process(name: str, cmd: list):
    """启动子进程并返回进程对象"""
    log(f'正在启动 {name} 服务...', "INFO")

    try:
        # 创建日志目录
        log_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logs')
        os.makedirs(log_dir, exist_ok=True)

        # 创建日志文件
        log_file = os.path.join(log_dir, f'{name.lower().replace(" ", "_")}_output.log')

        # 显示启动命令
        cmd_str = ' '.join(cmd)
        log(f"执行命令: {cmd_str}", "INFO")

        process = subprocess.Popen(
            cmd,
            cwd=os.path.dirname(os.path.abspath(__file__)),
            stdout=open(log_file, 'a', encoding='utf-8'),
            stderr=subprocess.STDOUT,
            text=True
        )

        # 等待一小段时间检查进程是否成功启动
        time.sleep(1)
        if process.poll() is None:
            log(f'[OK] {name} 服务启动成功 (PID: {process.pid})', "SUCCESS")
            log(f"  输出日志: {log_file}", "INFO")
        else:
            log(f'[FAIL] {name} 服务启动失败 (退出码: {process.returncode})', "ERROR")
            return None

        return process

    except Exception as e:
        log(f'[ERROR] {name} 服务启动异常: {str(e)}', "ERROR")
        raise

def start_fastapi():
    """启动FastAPI服务"""
    return start_process(
        'FastAPI',
        [
            sys.executable, '-m', 'uvicorn',
            'main:app',
            '--host', '0.0.0.0',
            '--port', '8000',
            '--reload'
        ]
    )

def start_celery_worker():
    """启动Celery Worker"""
    # 根据操作系统选择合适的池类型
    import platform
    if platform.system() == 'Windows':
        pool_type = 'solo'  # Windows环境推荐使用solo池
        log("检测到Windows环境，使用solo池", "INFO")
    else:
        pool_type = 'prefork'  # Linux/Unix环境使用prefork池，性能更好
        log("检测到Linux/Unix环境，使用prefork池", "INFO")

    return start_process(
        'Celery Worker',
        [
            sys.executable, '-m', 'celery',
            '-A', 'app.celery_app', 'worker',
            '--loglevel=info',  # 降低日志级别，减少输出
            '--queues=monitor,notification,cleanup',
            f'--pool={pool_type}',
            '--concurrency=2',  # 限制并发数
            '--without-heartbeat',  # 禁用心跳，减少网络问题
            '--without-gossip',     # 禁用gossip，减少网络通信
            '--without-mingle'      # 禁用mingle，减少启动时的网络通信
        ]
    )

def start_celery_beat():
    """启动Celery Beat定时任务调度器"""
    return start_process(
        'Celery Beat',
        [
            sys.executable, '-m', 'celery',
            '-A', 'app.celery_app', 'beat',
            '--loglevel=debug'
        ]
    )

def monitor_processes(processes):
    """监控进程状态"""
    while True:
        time.sleep(10)  # 每10秒检查一次
        for i, process in enumerate(processes):
            if process and process.poll() is not None:
                service_names = ['FastAPI', 'Celery Worker', 'Celery Beat']
                service_name = service_names[i] if i < len(service_names) else f'Service-{i}'
                log(f'[WARNING] {service_name} 服务异常退出 (退出码: {process.returncode})', "ERROR")

def print_service_status(processes):
    """打印服务状态"""
    log("当前服务状态:", "HEADER")
    service_names = ['FastAPI', 'Celery Worker', 'Celery Beat']

    for i, process in enumerate(processes):
        if process:
            service_name = service_names[i] if i < len(service_names) else f'Service-{i}'
            if process.poll() is None:
                log(f"  [OK] {service_name}: 运行中 (PID: {process.pid})", "SUCCESS")
            else:
                log(f"  [FAIL] {service_name}: 已停止 (退出码: {process.returncode})", "ERROR")
        else:
            service_name = service_names[i] if i < len(service_names) else f'Service-{i}'
            log(f"  [FAIL] {service_name}: 启动失败", "ERROR")

def main():
    # 打印启动横幅
    print_banner()

    # 系统检查
    log("开始系统检查...", "HEADER")

    if not check_dependencies():
        log("依赖检查失败，退出启动", "ERROR")
        return

    check_ports()

    log("系统检查完成，开始启动服务...", "SUCCESS")
    log("-" * 50, "INFO")

    processes = []
    try:
        # 启动所有服务
        log("正在启动核心服务...", "HEADER")

        fastapi_process = start_fastapi()
        if fastapi_process:
            processes.append(fastapi_process)

        worker_process = start_celery_worker()
        if worker_process:
            processes.append(worker_process)

        beat_process = start_celery_beat()
        if beat_process:
            processes.append(beat_process)

        # 检查是否有服务启动失败
        failed_services = [p for p in processes if p is None]
        if failed_services:
            log(f"有 {len(failed_services)} 个服务启动失败", "ERROR")
            return

        log("-" * 50, "INFO")
        log('[SUCCESS] 所有服务启动完成！', "SUCCESS")
        log('[INFO] FastAPI文档: http://localhost:8000/docs', "INFO")
        log('[INFO] 系统监控: 查看 logs/ 目录下的日志文件', "INFO")
        log('[NOTICE] 按 Ctrl+C 停止所有服务', "WARNING")
        log("-" * 50, "INFO")

        # 显示服务状态
        print_service_status(processes)

        # 等待所有进程完成
        for process in processes:
            if process:
                process.wait()

    except KeyboardInterrupt:
        log('', "INFO")  # 换行
        log('[STOP] 接收到停止信号，正在关闭所有服务...', "WARNING")

        for i, process in enumerate(processes):
            if process and process.poll() is None:  # 如果进程仍在运行
                service_names = ['FastAPI', 'Celery Worker', 'Celery Beat']
                service_name = service_names[i] if i < len(service_names) else f'Service-{i}'
                try:
                    if os.name == 'nt':  # Windows
                        process.terminate()
                    else:  # Unix/Linux
                        process.send_signal(signal.SIGTERM)

                    log(f'[OK] 已停止 {service_name} (PID: {process.pid})', "SUCCESS")
                except Exception as e:
                    log(f'[FAIL] 停止 {service_name} 失败 (PID: {process.pid}): {str(e)}', "ERROR")

    except Exception as e:
        log(f'系统运行异常: {str(e)}', "ERROR")

    finally:
        log('[FINISH] 所有服务已关闭，感谢使用！', "SUCCESS")

if __name__ == '__main__':
    main()